﻿using System;
using System.Collections.Generic;
using System.Drawing;
using System.Linq;
using System.Text;
using OpenCLNet;

namespace GPUTracer
{
    public class CPURayTracer
    {
        private static KDTreeElement _kdTree;
        private static Object[] _objects;
        private static int _numPlanes;
        private static Light[] _lights;
        private static Bitmap _skyTexture = new Bitmap("skyTexture.png");
        private static Bitmap _woodTexture = new Bitmap("woodTexture.JPG");

        private static readonly int _DEPTH = 3;
        private static readonly bool BRUTEFORCERENDERING = false;
        private static readonly bool SHOW_SHADOWS = true;

        public static void RayTrace(Bitmap OutputBitmap, KDTreeElement root, Object[] objects, int numPlanes, Light[] lights, Float4 cameraPosition, Float4 cameraForward, Float4 cameraRight, Float4 cameraUp)
        {
            _objects = objects;
            _numPlanes = numPlanes;
            _lights = lights;
            _kdTree = root;

            for (int y = 0; y < OutputBitmap.Height; y++)
            {
                for (int x = 0; x < OutputBitmap.Width; x++)
                {
                    Float4 rayDirection = getRayDirection(OutputBitmap.Width, OutputBitmap.Height, x, y, cameraForward, cameraRight, cameraUp);
                    Float4 pixelColor = traceRay(cameraPosition, rayDirection);
                    var c = Color.FromArgb((int)(pixelColor.S0 > 1 ? 255 : pixelColor.S0 * 255),
                                           (int)(pixelColor.S1 > 1 ? 255 : pixelColor.S1 * 255),
                                           (int)(pixelColor.S2 > 1 ? 255 : pixelColor.S2 * 255));
                    OutputBitmap.SetPixel(x, y, c);
                }
            }
        }

        private static Float4 traceRay(Float4 rayStart, Float4 rayDirection)
        {
            Float4 color = new Float4(0.0f,0.0f,0.0f,0.0f);
	        float reflect = 1.0f;
	        int depth = _DEPTH;

	        Object object2;
	        Float4 pos2;
	        Float4 normal2;
	        Float4 reflectDir2;

	        for (int i = 0; i < depth; i++)
	        {
		        float distance;
		        Object obj = intersectingObject(rayStart, rayDirection, out distance);
		        if (obj == null)
			        return color;

		        Float4 pos = (rayDirection.Times(distance)).Add(rayStart);

		        //shade:
		        Float4 normal = getNormal(obj, pos);
		        Float4 reflectDir = rayDirection.Sub((normal.Times(Helper.Dot(normal, rayDirection)).Times(2)));

                color = color.Add(getNaturalColor(pos, normal, reflectDir, obj).Times(reflect));

		        //reflection/refraction:
		        reflect *= getReflect(obj.Surface, pos);
		        if (reflect == 0.0f)
			        break;
		        if (reflect < 0.0f)
		        {
			        reflect = -reflect;

			        if ((i < depth-1) && (obj.Surface == Object.SurfaceType.TransparentAndShiny))	//for surface "transparent and shiny"
			        {
				        //compute the color of the reflection ray with depth 1 (to fake the ability to reflect and refract at the same time):
				        object2 = intersectingObject(pos, reflectDir, out distance);			
				        if (object2 != null)
				        {
					        pos2 = (reflectDir.Times(distance)).Add(pos);
					        normal2 = getNormal(object2, pos2);
					        reflectDir2 = reflectDir.Sub(normal2.Times(2 * Helper.Dot(normal2, reflectDir)));
                            color = color.Add(getNaturalColor(pos2, normal2, reflectDir2, object2).Times(reflect));
				        }
			        }

			        if (getRefractionRay(obj, pos, rayDirection, out rayStart, out rayDirection))
				        continue;
		        }
		        rayStart = pos;
		        rayDirection = reflectDir;
	        }
	        return color;
        }

        private static float getReflect(Object.SurfaceType surface, Float4 pos)
        {
            switch (surface)
            {
                case Object.SurfaceType.Checkerboard:	//checkerboard
                    return (int)(Math.Floor(pos.S2) + Math.Floor(pos.S0)) % 2 != 0 ? 0.05f : 0.1f;
                case Object.SurfaceType.Shiny: //shiny
                    return 0.7f;
                case Object.SurfaceType.Transparent: //transparent
                case Object.SurfaceType.TransparentAndShiny:	//transparent and shiny
                    return -0.9f;
                case Object.SurfaceType.Sky:	//sky
                case Object.SurfaceType.Wood:	//wood
                    return 0.0f;
            }
            return 0.0f;
        }

        private static Float4 getNaturalColor(Float4 pos, Float4 normal, Float4 rd, Object obj)
        {
	        bool isInShadow = false;

            Float4 ret = new Float4();

	        for (int i = 0; i < _lights.Length; i++)
	        {
		        Float4 ldis = _lights[i].Position.Sub(pos);
		        Float4 livec = ldis.Normalize();

                if (SHOW_SHADOWS)
                {
                    if (obj.Surface != Object.SurfaceType.Sky  && obj.Surface != Object.SurfaceType.Transparent 
                        && obj.Surface != Object.SurfaceType.TransparentAndShiny) //sky and transparent don't receive shadows
                    {
                        float neatIsect = testRay(pos, livec);
                        isInShadow = !((neatIsect > ldis.Length()) || (neatIsect == 0.0f));
                    }
                }

	            if (!isInShadow)
                {
                    float illum = Helper.Dot(livec, normal);
			        Float4 lcolor;
			        if(illum > 0){
				        lcolor= _lights[i].Color.Times(illum);
			        }else{
				        lcolor=new Float4();
			        }

			        float roughness = getRoughness(obj.Surface);

			        float spec = Helper.Dot(livec, rd.Normalize());
			        Float4 scolor;
			        if (spec > 0.0f) {
				        scolor = _lights[i].Color.Times((float) Math.Pow(spec, roughness));
			        }else{
				        scolor = new Float4();
			        }
			
			        Float4 diffuse = getDiffuse(obj, pos);
			        Float4 specular = getSpecular(obj.Surface, pos);

                    ret = ret.Add(new Float4(lcolor.S0*diffuse.S0, lcolor.S1*diffuse.S1, lcolor.S2*diffuse.S2, 0.0f)).Add
					          (new Float4(scolor.S0*specular.S0, scolor.S1*specular.S1, scolor.S2*specular.S2, 0.0f));
                }
	        }

	        return ret;
        }

        private static float testRay(Float4 rayStart, Float4 rayDirection)
        {
            float distance;
            if (intersectingObject(rayStart, rayDirection, out distance) == null)
                return 0.0f;

            return distance;
        }

        private static float getRoughness(Object.SurfaceType surface)
        {
            switch (surface)
            {
                case Object.SurfaceType.Checkerboard:	//checkerboard
                    return 150.0f;
                case Object.SurfaceType.Shiny:          //shiny
                    return 50.0f;
                case Object.SurfaceType.Transparent:    //transparent
                case Object.SurfaceType.TransparentAndShiny:	//transparent and shiny
                    return 10.0f;
                case Object.SurfaceType.Sky:	        //sky
                case Object.SurfaceType.Wood:	        //wood
                    return 0.0f;
            }
            return 0.0f;
        }

        private static Float4 getDiffuse(Object o, Float4 pos)
        {
	        float u, v;

	        switch (o.Surface)
	        {
		        case Object.SurfaceType.Checkerboard:	//checkerboard
			        return (int)(Math.Floor(pos.S2) + Math.Floor(pos.S0)) % 2 != 0 ? new Float4(1.0f,1.0f,1.0f,0.0f) : new Float4(0.0f,0.0f,0.0f,0.0f);
		        case Object.SurfaceType.Shiny: //shiny
			        return new Float4(1.0f, 1.0f, 1.0f, 0.0f);
		        case Object.SurfaceType.Transparent: //transparent
		        case Object.SurfaceType.TransparentAndShiny:	//transparent and shiny
			        return new Float4(1.0f, 1.0f, 1.0f, 0.0f);
		        case Object.SurfaceType.Sky:	//sky
		        case Object.SurfaceType.Wood:	//wood
	                v = (float) Math.Sqrt(Math.Pow(pos.S0, 2) + Math.Pow(pos.S1, 2) + Math.Pow(pos.S2, 2));
			        u = pos.S0 / v;
			        v = pos.S1 / v;

	                u = (float) (u - Math.Truncate(u));
	                v = (float) (v - Math.Truncate(v));

                    if (u < 0)
                        u = 1.0f + u;
                    if (v < 0)
                        v = 1.0f + v;

	                Color pix;
	                Bitmap tex;
                    if (o.Surface == Object.SurfaceType.Sky)
                        tex = _skyTexture;
                    else
                        tex = _woodTexture;

                    pix = tex.GetPixel((int)(u * tex.Width) % tex.Width, (int)(v * tex.Height) % tex.Height);
			        return new Float4(pix.R/255.0f, pix.G/255.0f, pix.B/255.0f, 0.0f);		
	        }
	        return new Float4(0.0f, 0.0f, 0.0f, 0.0f);
        }

        private static Float4 getSpecular(Object.SurfaceType surface, Float4 pos)
        {
            switch (surface)
	        {
		        case Object.SurfaceType.Checkerboard:	//checkerboard
			        return new Float4(1.0f, 1.0f, 1.0f, 0.0f);
		        case Object.SurfaceType.Shiny: //shiny
			        return new Float4(0.5f, 0.5f, 0.5f, 0.0f);
		        case Object.SurfaceType.Transparent: //transparent
		        case Object.SurfaceType.TransparentAndShiny:	//transparent and shiny
			        return new Float4(0.5f, 0.5f, 0.5f, 0.0f);
		        case Object.SurfaceType.Sky:	//sky
		        case Object.SurfaceType.Wood:	//wood
			        return new Float4(0.0f, 0.0f, 0.0f, 0.0f);
	        }
            return new Float4(0.0f, 0.0f, 0.0f, 0.0f);
        }

        private static bool getRefractionRay(Object o, Float4 pos, Float4 rayDirection, out Float4 outStart, out Float4 outDir)
        {
            outStart = new Float4();
            outDir = new Float4();

            if ((o.Surface != Object.SurfaceType.Transparent) && o.Surface != Object.SurfaceType.TransparentAndShiny)
		        return false;

	        Float4 normal = getNormal(o, pos);
	        Float4 refractDir;
	        Float4 VplusN;
	        float r;
	        float kf;
	        float k;

	        if (o is Sphere)
	        {
	            k = 1.01f; //index of refraction

	            r = -Helper.Dot(normal, rayDirection);
	            VplusN = (normal.Add(rayDirection.Times(1/r)));
	            kf = (float) (Math.Pow(k, 2)*Math.Pow(1.0f/r, 2) - Helper.Dot(VplusN, VplusN));
	            if (kf < 0.0f) //total internal reflection
	            {
	                return false;
	            }
	            kf = (float) (1.0f/Math.Sqrt(kf));
                refractDir = normal.Times(-1).Add(VplusN.Times(kf)).Normalize();
	            r = -Helper.Dot(normal, refractDir);
	            outStart = pos.Add(refractDir.Times(r*(2* ((Sphere)o).Radius)));

	            //*outStart is now the point where the ray comes out of the sphere again.
	            //the direction inside is refractDir. However, by leaving the sphere medium, the direction has to be modified again:
	            normal = getNormal(o, outStart).Times(-1);
	            r = -Helper.Dot(normal, refractDir);
	            VplusN = (normal.Add(refractDir.Times(1/r)));
	            kf = (float) (Math.Pow(1.0f/k, 2)*Math.Pow(1.0f/r, 2) - Helper.Dot(VplusN, VplusN));
	            if (kf < 0.0f) //total internal reflection
	            {
	                return false;
	            }
	            kf = (float) (1.0f/Math.Sqrt(kf));
                outDir = normal.Times(-1).Add(VplusN.Times(kf)).Normalize();
	        }
            else if (o is Plane)
            {
			    k = 1.11f;		//index of refraction

			    r = -Helper.Dot(normal, rayDirection);
			    VplusN = (normal.Add(rayDirection.Times(1/r)));
			    kf = (float) (Math.Pow(k, 2) * Math.Pow(1.0f/r, 2) - Helper.Dot(VplusN, VplusN));
			    if (kf < 0.0f)		//total internal reflection
			    {
				    return false;
			    }
			    kf = (float) (1.0f/Math.Sqrt(kf));
			    outDir = (normal.Add(VplusN.Times(kf))).Times(-1).Normalize();
			    outStart = pos.Sub(normal.Times(0.03f));
            }

	        return true;
        }

        private static Float4 getNormal(Object o, Float4 pos)
        {
            if (o is Sphere)
            {
                return pos.Sub(((Sphere) o).Position).Normalize();
            }
            if (o is Plane)
            {
                return ((Plane) o).Normal.Normalize();
            }
            if (o is Triangle)
            {
                return Helper.Cross(((Triangle) o).P2.Sub(((Triangle) o).P1), ((Triangle) o).P3.Sub(((Triangle) o).P1)).Normalize();
            }
            return new Float4();
        }

        private static Object intersectingObject(Float4 rayStart, Float4 rayDirection, out float distanceOut)
        {
	        Object intersectingObject = null;
            distanceOut = float.MaxValue;

            if (BRUTEFORCERENDERING)
            {
                return intersectingObjectInNode(rayStart, rayDirection, _objects, out distanceOut);
            }

	        float planeDistanceOut;
	        Object intersectingPlane = intersectingObjectInNode(rayStart, rayDirection, _objects, _numPlanes, out planeDistanceOut);
	
	        Float4 bboxMin = _kdTree.BBox.Min;
	        Float4 bboxMax = _kdTree.BBox.Max;
	        
	        KDTreeElement kdcurrent = _kdTree;
            
            Stack<Subroom> kdStack = new Stack<Subroom>();

	        float global_tmin;
	        float global_tmax;
	        float tmin;
	        float tmax;
	        float thit;

	        KDTreeElement first;
            KDTreeElement second;

            float rayAxisDirection, rayAxisStart;
	
	        if (!intersectBBox(rayStart, rayDirection, bboxMin, bboxMax, out global_tmin, out global_tmax))
	        {
		        distanceOut = planeDistanceOut;
		        return intersectingPlane;
	        }

	        tmin = global_tmin;
	        tmax = global_tmax;

	        do
	        {
		        if (kdcurrent is KDTreeLeaf)
		        {
			        if (kdcurrent.NumObjects != 0)
			        {
				        intersectingObject = intersectingObjectInNode(rayStart, rayDirection, 
															          ((KDTreeLeaf)kdcurrent).ContainedObjects.ToArray(), out distanceOut);
			        }
			        if (intersectingObject == null)
			        {
				        if (kdStack.Count == 0)
				            break;

			            var top = kdStack.Pop();
			            kdcurrent = top.current;
			            tmin = top.tmin;
			            tmax = top.tmax;
			        }
		        }
		        else										//is node
		        {
			        switch (((KDTreeNode)kdcurrent).SplitAxis)
			        {
				        case Axis.XAxis:	//X Axis
					        rayAxisDirection = rayDirection.S0;
					        rayAxisStart = rayStart.S0;
					        break;
				        case Axis.YAxis:	//Y Axis
					        rayAxisDirection = rayDirection.S1;
					        rayAxisStart = rayStart.S1;
					        break;
				        case Axis.ZAxis:	//Z Axis
                        default:
					        rayAxisDirection = rayDirection.S2;
					        rayAxisStart = rayStart.S2;
					        break;
			        }

			        thit = (((KDTreeNode)kdcurrent).SplitPosition - rayAxisStart) / rayAxisDirection;

			        if (rayAxisStart < ((KDTreeNode)kdcurrent).SplitPosition)
			        {
			            first = ((KDTreeNode) kdcurrent).LeftChild;
                        second = ((KDTreeNode)kdcurrent).RightChild;
			        }
			        else
			        {
                        first = ((KDTreeNode)kdcurrent).RightChild;
                        second = ((KDTreeNode)kdcurrent).LeftChild;
			        }

                    if ((thit >= tmax) || (thit < 0))
			        {
                        kdcurrent = first;
			        }
			        else if (thit <= tmin)
			        {
				        kdcurrent = second;
			        }
			        else
			        {
				        //push second node on stack:
				        kdStack.Push(new Subroom() {current =  second, tmax = tmax, tmin = thit});

				        //current node is the first one:
				        tmax = thit;
				        kdcurrent = first;
			        }
		        }
	        } while (intersectingObject == null);
            
            if ((intersectingObject == null) || (intersectingPlane != null && (distanceOut > planeDistanceOut)))
	        {
		        intersectingObject = intersectingPlane;
		        distanceOut = planeDistanceOut;
	        }

	        return intersectingObject;
        }

        private static bool intersectBBox(Float4 rayStart, Float4 rayDirection, Float4 bboxMin, Float4 bboxMax, out float tmin, out float tmax)
        {
            tmin = 0.0f;
            tmax = 0.0f;
            float t0 = float.MinValue;
            float t1 = float.MaxValue;

            float invRayDir = 1.0f / rayDirection.S0;
            float tNear = (bboxMin.S0 - rayStart.S0) * invRayDir;
            float tFar = (bboxMax.S0 - rayStart.S0) * invRayDir;

            if (tNear > tFar)
            {
                invRayDir = tFar;
                tFar = tNear;
                tNear = invRayDir;
            }
            t0 = tNear > t0 ? tNear : t0;
            t1 = tFar < t1 ? tFar : t1;
            if (t0 > t1)
                return false;

            invRayDir = 1.0f / rayDirection.S1;
            tNear = (bboxMin.S1 - rayStart.S1) * invRayDir;
            tFar = (bboxMax.S1 - rayStart.S1) * invRayDir;

            if (tNear > tFar)
            {
                invRayDir = tFar;
                tFar = tNear;
                tNear = invRayDir;
            }
            t0 = tNear > t0 ? tNear : t0;
            t1 = tFar < t1 ? tFar : t1;
            if (t0 > t1)
                return false;

            invRayDir = 1.0f / rayDirection.S2;
            tNear = (bboxMin.S2 - rayStart.S2) * invRayDir;
            tFar = (bboxMax.S2 - rayStart.S2) * invRayDir;

            if (tNear > tFar)
            {
                invRayDir = tFar;
                tFar = tNear;
                tNear = invRayDir;
            }
            t0 = tNear > t0 ? tNear : t0;
            t1 = tFar < t1 ? tFar : t1;
            if (t0 > t1)
                return false;

            tmin = t0;
            tmax = t1;

            return true;
        }

        private static Object intersectingObjectInNode(Float4 rayStart, Float4 rayDirection, Object[] nodeObjects, int numObjects, out float distanceOut)
        {
            float minDistance = float.MaxValue;
            Object minObject = null;

            float distance = float.MaxValue;
            float t;

            for (int i = 0; i < numObjects; i++)
            {
                if (nodeObjects[i] is Sphere)
                {
                    intersectingSphere(((Sphere)nodeObjects[i]).Position, ((Sphere)nodeObjects[i]).Radius,
                                       rayStart, rayDirection, out distance);
                }
                else if (nodeObjects[i] is Plane)
                {
                    t = Helper.Dot(((Plane)nodeObjects[i]).Normal.Normalize(), rayDirection);
                    if (t <= 0.0f)
                    {
                        distance = (Helper.Dot(((Plane)nodeObjects[i]).Normal.Normalize(), rayStart) +
                                    ((Plane)nodeObjects[i]).Offset) / (-t);
                    }
                }
                else if (nodeObjects[i] is Triangle)
                {
                    intersectingTriangle(((Triangle)nodeObjects[i]).P1, ((Triangle)nodeObjects[i]).P2,
                                         ((Triangle)nodeObjects[i]).P3, rayStart, rayDirection, out distance);
                }

                if (distance < minDistance && distance > 0)
                {
                    minObject = nodeObjects[i];
                    minDistance = distance;
                }
            }
            if (minDistance == float.MaxValue)
            {
                distanceOut = 0f;
                return null;
            }

            distanceOut = minDistance;
            return minObject;
        }

        private static Object intersectingObjectInNode(Float4 rayStart, Float4 rayDirection, Object[] nodeObjects, out float distanceOut)
        {
            return intersectingObjectInNode(rayStart, rayDirection, nodeObjects, nodeObjects.Count(), out distanceOut);
        }

        private static bool intersectingTriangle(Float4 p1, Float4 p2, Float4 p3, Float4 rayStart, Float4 rayDirection, out float distanceOut)
        {
            distanceOut = 0.0f;
            Float4 e1 = p2.Sub(p1);
            Float4 e2 = p3.Sub(p1);
            Float4 s1 = Helper.Cross(rayDirection, e2);
            float divisor = Helper.Dot(s1, e1);
            if (divisor == 0.0f)
                return false;
            float invDivisor = 1.0f / divisor;

            Float4 d = rayStart.Sub(p1);
            float b1 = Helper.Dot(d, s1) * invDivisor;
            if ((b1 < 0.0f) || (b1 > 1.0f))
                return false;

            Float4 s2 = Helper.Cross(d, e1);
            float b2 = Helper.Dot(rayDirection, s2) * invDivisor;
            if ((b2 < 0.0f) || (b1 + b2 > 1.0f))
                return false;

            distanceOut = Helper.Dot(e2, s2) * invDivisor;
            return true;
        }

        private static bool intersectingSphere(Float4 pos, float radius, Float4 rayStart, Float4 rayDirection, out float distanceOut)
        {
            float denom;
            float disc;
            Float4 eo;

            eo = pos.Sub(rayStart);
            denom = Helper.Dot(eo, rayDirection);
            if (denom >= 0)
            {
                disc = (float) (Math.Pow(radius, 2) - (Helper.Dot(eo, eo) - Math.Pow(denom, 2)));
                if (disc >= 0.0f)
                {
                    distanceOut = (float) (denom - Math.Sqrt(disc));
                    return true;
                }
            }
            distanceOut = 0f;
            return false;
        }

        private struct Subroom
        {
            public float tmin;
            public float tmax;
            public KDTreeElement current;
        }

        private static Float4 getRayDirection(int width, int height, int x, int y, Float4 cameraForward, Float4 cameraRight, Float4 cameraUp)
        {
            float ratio = (float)width / height;
            float recenteredX = (x - (width / 2)) / (2.0f * width) * ratio;
            float recenteredY = (y - (height / 2)) / (2.0f * height);
            return cameraForward.Add(cameraRight.Times(recenteredX)).Add(cameraUp.Times(recenteredY)).Normalize();
        }
    }
}
